package com.zj.eswrapper;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import com.zj.eswrapper.interfaces.impl.AbstractEsQueryWrapper;
import com.zj.eswrapper.wrapper.EsQueryWrapper;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.unit.DistanceUnit;
import org.elasticsearch.index.query.*;

import java.math.BigDecimal;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author arthur_zhou
 * @date 2024/1/29 15:35
 */
public class MiddlewareQueryBuilderImpl implements IMiddlewareQueryBuilder<QueryBuilder> {

    @Override
    public AbstractEsQueryWrapper<String, ?, ?, ?, QueryBuilder, ?, ?> createWrapper() {
        return new EsQueryWrapper<>();
    }

    @Override
    public void multiMatch(AbstractEsQueryWrapper<String, ?, ?, ?, QueryBuilder, ?, ?> wrapper, Key key) {
        if (Objects.isNull(key)) {
            return;
        }
        String query = key.getQuery();
        if (StrUtil.isEmpty(query)) {
            return;
        }
        MultiMatchQueryBuilder builder = QueryBuilders.multiMatchQuery(query);
        List<Key.Field> fields = key.getFields();
        if (CollectionUtil.isEmpty(fields)) {
            return;
        }
        for (Key.Field field : fields) {
            if (StrUtil.isEmpty(field.getField())) {
                continue;
            }
            builder.field(field.getField(), field.getWeight());
        }
        String defaultOperator = key.getDefaultOperator();
        // 默认按照AND计算
        if ("OR".equalsIgnoreCase(defaultOperator)) {
            builder.operator(Operator.OR);
        } else {
            builder.operator(Operator.AND);
        }
        wrapper.must().custom(builder);
    }

    @Override
    public void graph(AbstractEsQueryWrapper<String, ?, ?, ?, QueryBuilder, ?, ?> wrapper, Graph graph) {
        String field = graph.getField();
        if (StrUtil.isEmpty(field)) {
            return;
        }
        List<Graph.Point> points = graph.getPoints();
        if (CollectionUtil.isEmpty(points)) {
            return;
        }
        List<GeoPoint> pointList = graph.getPoints().stream().map(v -> new GeoPoint(v.getLat(), v.getLon())).collect(Collectors.toList());
        GeoPolygonQueryBuilder builder = QueryBuilders.geoPolygonQuery(field, pointList);
        wrapper.custom(builder);
    }

    @Override
    public void geoDistance(AbstractEsQueryWrapper<String, ?, ?, ?, QueryBuilder, ?, ?> wrapper, String key, BigDecimal radius, BigDecimal lon, BigDecimal lat) {
        if (Objects.isNull(radius) || Objects.isNull(lon) || Objects.isNull(lat)) {
            return;
        }
        wrapper.custom(() -> {
            GeoDistanceQueryBuilder builder = QueryBuilders.geoDistanceQuery(key);
            builder.point(lat.doubleValue(), lon.doubleValue());
            builder.distance(radius.doubleValue(), DistanceUnit.METERS);
            return builder;
        });
    }
}
